#!/usr/bin/python
# -*- coding: UTF-8 -*-

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
# 系统库
import os
from math import exp
import numpy as np
# 框架
from PIL import Image
import cv2
from model_service.pytorch_model_service import PTServingBaseService
import torch.nn as nn
import torch
import logging
import torchvision.models as models
import torchvision.transforms as transforms
# 自定义模型
from resnet import resnet50,resnext50_32x4d,resnext101_32x8d

# 图像增强代码
infer_transformation = transforms.Compose([
    transforms.Resize((288,288)),
    transforms.RandomHorizontalFlip(), # 依概率p水平翻转
    transforms.ToTensor(),
    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])


normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
ttatransform = transforms.Compose([
    transforms.Resize(256),
    transforms.TenCrop(244), # this is a list of PIL Images
    transforms.Lambda(lambda crops: torch.stack([normalize(transforms.ToTensor()(crop)) for crop in crops])), # returns a 4D tensor
])


IMAGES_KEY = 'images'
MODEL_INPUT_KEY = 'images'
MODEL_OUTPUT_KEY = 'logits'
LABELS_FILE_NAME = 'labels_10c.txt'

class TTAModule(nn.Module):
    def __init__(self, model: nn.Module):
        super().__init__()
        self.model = model

    def tta_out(self, out):
        _, preds = torch.max(out.data, 1)
        li = preds.tolist()
        smax = sid = -1
        for i in li:
            if li.count(i) > smax:
                smax = li.count(i)
                sid = i
        rid = li.index(sid)
        return out[rid].unsqueeze(0)
    
    def forward(self, x):
        x = x.squeeze(0)
        result = self.model(x)
        result = self.tta_out(result)

        return result

def decode_image(file_content):

    """
    将字节解码为单个图像
    :param file_content: bytes
    :return: ndarray with rank=3
    """

    image = Image.open(file_content)
    image = image.convert('RGB')
    # print(image.shape)
   # image = np.asarray(image, dtype=np.float32)
    return image
#    image_content = r.files[file_content].read() # python 'List' class that holds byte
#    np_array = np.fromstring(image_content, np.uint8) # numpy array with dtype np.unit8
#    img_array = cv2.imdecode(np_array, cv2.IMREAD_COLOR) # numpy array in shape [height, width, channels]
 

def read_label_list(path):

    """
    从路径读取标签列表
    :param path: a path
    :return: a list of label names like: ['label_a', 'label_b', ...]
    """
    with open(path, 'r',encoding="utf8") as f:
        label_list = f.read().split(os.linesep)
    label_list = [x.strip() for x in label_list if x.strip()]
    # print(' label_list',label_list)
    return label_list


class FoodPredictService(PTServingBaseService):
    def __init__(self, model_name, model_path):

        global LABEL_LIST
        super(FoodPredictService, self).__init__(model_name, model_path)
        print(model_name, model_path)
        # 初始化模型
        self.model = set_model(model_path)
        dir_path = os.path.dirname(os.path.realpath(self.model_path))
        # 从路径读取标签列表
        LABEL_LIST = read_label_list(os.path.join(dir_path, LABELS_FILE_NAME))

    # 预处理
    def _preprocess(self, data):

        """
        Upredict服务根据输入数据提供“数据”。 就像这样：
          {
              'images': {
                'image_a.jpg': b'xxx'
              }
          }
        现在，一次预测一个图像。
        """
        preprocessed_data = {}
        input_batch = []
        for file_name, file_content in data[IMAGES_KEY].items():
            
            print('\tAppending image: %s' % file_name)

            image1 = decode_image(file_content)
            if torch.cuda.is_available():
                input_batch.append(infer_transformation(image1).cuda())
            else:
                input_batch.append(infer_transformation(image1))
        input_batch_var = torch.autograd.Variable(torch.stack(input_batch, dim=0), volatile=True)
        preprocessed_data[MODEL_INPUT_KEY] = input_batch_var

        # print('preprocessed_data',input_batch_var.shape())

        return preprocessed_data

    # 后期过程
    def _postprocess(self, data):

        """
        数据是模型的结果。 就像这样：
          {
            'logits': [[0.1, -0.12, 0.72, ...]]
          }
        logits的值是列表的单个列表，因为现在一次只能预测一张图像。
        """

        # logits_list = [0.1, -0.12, 0.72, ...]

        logits_list =  data['images'][0].detach().numpy().tolist()
        maxlist=max(logits_list)
        z_exp = [exp(i-maxlist) for i in  logits_list]
        
        sum_z_exp = sum(z_exp)
        softmax = [round(i / sum_z_exp, 3) for i in z_exp]

        # labels_to_logits = {
        #     'label_a': 0.1, 'label_b': -0.12, 'label_c': 0.72, ...
        # }

        labels_to_logits = {
            LABEL_LIST[i]: s for i, s in enumerate(softmax)
            # LABEL_LIST[i]: s for i, s in enumerate(logits_list)
        }

        predict_result = {
            MODEL_OUTPUT_KEY: labels_to_logits
        }

        return predict_result


def set_model(model_path, **kwargs):

    """Constructs a ResNet-50 model.
    Args:
        pretrained (bool): If True, returns a model pre-trained on ImageNet
    """
    model = resnext101_32x8d(pretrained=False)
    num_ftrs = model.fc.in_features
    model.fc = nn.Sequential(
        nn.Linear(num_ftrs, 10),
    )
    
    model.load_state_dict(torch.load(model_path, map_location ='cpu'))
    # save_model = torch.load(model_path, map_location ='cpu')
    # model.load_state_dict(save_model)
    
    # model = TTAModule(model)
    # model.load_state_dict(torch.load(model_path))

    model.eval()

    return model
